# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import time
from typing import List
from tortoise.queryset import Q
from tortoise.transactions import in_transaction
from hypertext import PagingResult
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.tools import ToolsUtil
from common.enums.agent import AgentEnum
from common.models.users import UserModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentRobotModel
from common.models.agents import AgentKnowArchiveModel
from common.models.agents import AgentChatRecordsModel
from common.postgres.public import PgKnowledgeModel, PgAttachmentModel
from apps.api.schemas.agent.robot import journal_schema as schema
from apps.api.service.agent.robot.teamed_service import RobotTeamedService


class RobotJournalService:

    @classmethod
    async def records(cls, params: schema.RobotJournalSearchIn, user_id: int) \
            -> PagingResult[schema.RobotJournalListVo]:
        """
        对话记录列表。

        Args:
            params (schema.RobotJournalSearchIn): 记录查询参数。
            user_id (int): 用户ID。

        Returns:
            PagingResult[schema.RobotJournalListVo]: 记录分页列表Vo。

        Author:
            zero
        """
        # 验证智能体
        robot = await AgentRobotModel.filter(code=params.bot, is_delete=0).first()
        if robot is None:
            raise AppException(str("智能体不存在!"))
        if robot.is_disable:
            raise AppException(str("智能体被禁用!"))

        # 验证权限
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, AgentEnum.POWER_VIEW)

        _where = [Q(robot_id=robot.id), Q(is_delete=0)]
        _where = _where + AgentChatRecordsModel.build_search({
            "=": ["scene"],
            "datetime": ["start_time|end_time@create_time"]
        }, params.model_dump())

        # 查询数据
        _model = AgentChatRecordsModel.filter(*_where).order_by("-id")
        _pager = await AgentChatRecordsModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size
        )

        # 查询用户
        user_ids = [item["user_id"] for item in _pager.lists if item["user_id"]]
        users = await UserModel.queryUserDictsById(user_ids)

        role_dict = {"system": "System", "user": "Human", "assistant": "AI"}

        # 处理数据
        _lists = []
        for item in _pager.lists:
            # 处理上下文
            ctx_array = []
            context = json.loads(item["context"] or "[]")
            for ctx in context:
                content = ctx["content"]
                if isinstance(ctx["content"], list):
                    annex, value = ("", "")
                    for ctx_item in ctx["content"]:
                        if ctx_item["type"] == "image_url":
                            annex += "Input an image\n"
                        elif ctx_item["type"] == "video_url":
                            annex += "Input an video\n"
                        elif ctx_item["type"] == "audio_url":
                            annex += "Input an audio\n"
                        elif ctx_item["type"] == "text":
                            value += ctx_item["text"] + "\n"
                    content = (annex + value).strip("\n")
                ctx_array.append({
                    "role": role_dict.get(ctx["role"], ctx["role"]),
                    "content": content
                })

            annex = json.loads(item["annex"] or "[]")
            for t in annex:
                t["file_path"] = await UrlUtil.to_absolute_url(t["file_path"])

            user = users.get(item["user_id"], {})
            if item["publish_id"]:
                user = {
                    "sn": item["publish_uid"],
                    "avatar": "static/images/default_avatar.png",
                    "nickname": "Guest"
                }

            _lists.append(schema.RobotJournalListVo(
                id=item["id"],
                scene=item["scene"],
                question=item["question"],
                answer=item["answer"],
                thinks=item["thinks"],
                annex=annex,
                quotes=json.loads(item["quotes"] or "[]"),
                context=ctx_array,
                use_tokens=item["use_tokens"],
                use_points=item["use_points"],
                user_delete=item["user_delete"],
                task_time=item["task_time"] + "s",
                create_time=item["create_time"],
                ip=item["ip"],
                user={
                    "sn": user.get("sn", ""),
                    "nickname": user.get("nickname", ""),
                    "avatar": await UrlUtil.to_absolute_url(user.get('avatar', ""))
                }
            ))

        _pager.lists = _lists
        return _pager

    @classmethod
    async def corrects(cls, user_id: int, post: schema.RobotJournalCorrectsIn):
        """
        对话记录修正。

        Args:
            user_id (int): 用户ID。
            post (schema.RobotJournalCorrectsIn): 修正参数。

        Author:
            zero
        """
        # 验智能体
        robot = await AgentRobotModel.filter(code=post.bot).first()
        if not robot:
            raise AppException("智能体不存在了")
        if robot.is_delete:
            raise AppException("智能体已被删除")
        if robot.is_disable:
            raise AppException("智能体已被禁用")

        # 验证权限
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, AgentEnum.POWER_ADMIN)

        # 验知识库
        know = await AgentKnowModel.filter(code=str(post.kb)).first()
        if not know:
            raise AppException("知识库不存在了")
        if know.is_delete:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")
        if know.user_id != user_id:
            raise AppException("知识库检测异常")

        # 验证文档
        archive = await AgentKnowArchiveModel.filter(code=post.fid, know_id=know.id).first()
        if not archive:
            raise AppException("文档不存在哦")
        if archive.is_delete:
            raise AppException("文档已被删除")

        # 数据查重
        text: str = post.question.strip() + post.answer.strip()
        signed: str = ToolsUtil.make_md5_str(text)
        pg = await (PgKnowledgeModel
                    .filter(signed=signed)
                    .filter(know_id=know.id)
                    .filter(archive_id=archive.id)
                    .filter(is_delete=0)
                    .first())

        if pg:
            raise AppException("检测到已存在重复数据,请勿重复写入。")

        async with in_transaction("pgsql"):
            # 插入数据
            result = await PgKnowledgeModel.create(
                know_id=know.id,
                archive_id=archive.id,
                import_uid=user_id,
                signed=signed,
                question=post.question,
                answer=post.answer,
                create_time=int(time.time()),
                update_time=int(time.time())
            )

            # 更新分词
            await PgKnowledgeModel.update_ts_vectory([str(result.uuid)])

            # 处理附件
            annex = []
            for item in post.annex:
                texts = (item.question or item.name).strip()
                signed = ToolsUtil.make_md5_str(texts)
                annex.append(PgAttachmentModel(
                    scene="know",
                    dataset_id=str(result.uuid),
                    archive_id=archive.id,
                    know_id=know.id,
                    type=item.type,
                    name=item.name,
                    path=item.path,
                    ext=item.ext,
                    size=item.size,
                    sort=item.sort,
                    signed=signed,
                    question=str(texts),
                    create_time=int(time.time()),
                    update_time=int(time.time())
                ))

            if annex:
                await PgAttachmentModel.bulk_create(annex)
                await PgAttachmentModel.update_ts_vectory(str(result.uuid))

    @classmethod
    async def delete(cls, user_id: int, bot: str, ids: List[int]) -> None:
        """
        对话记录删除。

        Args:
            bot (str): 智能体编号。
            ids (List[int]str): 记录ID。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 验证智能体
        robot = await AgentRobotModel.filter(code=bot, is_delete=0).first()
        if not robot:
            raise AppException(str("智能体不存在!"))
        if robot.is_disable:
            raise AppException(str("智能体被禁用!"))

        # 验证权限
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, AgentEnum.POWER_EDIT)

        # 查询记录
        records = await (AgentChatRecordsModel
                         .filter(id__in=ids)
                         .filter(robot_id=robot.id)
                         .filter(is_delete=0)
                         .all().values("id"))

        # 提取ID
        records_ids = [item["id"] for item in records]
        if not records_ids:
            raise AppException("无可删除数据!")

        # 删除数据
        await AgentChatRecordsModel.filter(id__in=records_ids).update(
            user_delete=1,
            delete_time=int(time.time())
        )
